# 新建 scripts/precompute_gallery.py  预计算所有图库特征
import sys
import os
from pathlib import Path


# 获取当前文件的绝对路径
current_file = Path(__file__).resolve()
# 计算项目根目录：上溯两级（假设train.py在scripts/目录下）
project_root = current_file.parent.parent
# 将项目根目录添加到系统路径
sys.path.insert(0, str(project_root))


import torch
from torchvision import transforms
from src.datasets.veri776_dataset import Veri776Dataset
from src.models.vehicle_transformer import VehicleTransformer
from tqdm import tqdm
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                            QHBoxLayout, QPushButton, QLabel, QFileDialog)
from PyQt5.QtGui import QPixmap, QImage
from PyQt5.QtCore import Qt, QThread, pyqtSignal
import numpy as np
import time
from PIL import Image
import argparse
from torch.utils.data import Dataset
import torch.nn as nn
import torch.nn.functional as F
from torchvision.models import vit_b_16
from einops import rearrange, repeat
from collections import defaultdict
from torch.utils.data import DataLoader  # 新增导入
from torch.optim import AdamW
from pathlib import Path
import matplotlib.pyplot as plt
from collections import Counter
from sklearn.metrics import average_precision_score  # 新增导入
from torch.optim.lr_scheduler import CosineAnnealingLR
import warnings
from torch.optim.lr_scheduler import LambdaLR


def extract_gallery_features():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # 加载模型，参数与训练配置完全一致
    model = VehicleTransformer(num_classes=776,
     img_size=(224, 224),
     patch_sizes=[16, 8], 
     local_parts=7,
     embed_dim=128,
     depth=4,
     num_heads=4,
     mlp_ratio=4).to(device)
    
    model.load_state_dict(torch.load("checkpoints/best_model.pth", weights_only=False)["state_dict"]) #加载最佳模型
    model.eval()
    


    # 加载图库数据
    transform = transforms.Compose([
         transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
      
    gallery_set = Veri776Dataset(mode="test", transform=transform) #加载测试集 ， mode="test" 指定图库数据
    gallery_camera_ids = gallery_set.camera_ids  # 直接获取数据集已解析的摄像头ID列表


    # 提取特征
    features, paths = [], []
    for sample in tqdm(gallery_set):
        img = sample["image"]
        path = sample["path"]


        with torch.no_grad():
            feat = model(img.unsqueeze(0).to(device))["bn_feature"] #提取BN Neck特征   使用归一化后的特征
            feat = F.normalize(feat, p=2, dim=1)  # 新增归一化  ，特征L2归一化 ，使余弦相似度等价于内积运算

            features.append(feat.cpu())
            paths.append(path)
    
    # 保存结果
    torch.save({
        "features": torch.cat(features),
        "paths": paths,
         "camera_ids": gallery_camera_ids  # 新增字段
    }, "gallery_features.pt")       #保存特征和路径   ，进制格式高效存储



if __name__ == "__main__":
    extract_gallery_features()




#核心功能：
#预计算图库特征：将测试集（gallery）所有车辆图像的特征提前提取并存储，加速后续检索过程
#输出文件：生成 gallery_features.pt 包含特征向量和对应图像路径

#加载训练好的模型：从checkpoints/best_model.pth加载最佳车辆重识别模型
#提取图库特征：对 VeRi-776 数据集中的测试集（图库）进行特征提取
#特征处理与存储：归一化特征并保存为二进制文件，包含特征向量、图像路径和摄像头 ID


#预计算后，检索流程可简化为：
#加载预存的gallery_features.pt
#计算查询图像特征
#计算查询特征与图库特征的相似度矩阵
#按相似度排序并返回结果
#相比未预计算的流程，省去了图库特征的重复计算步骤，大幅提升响应速度。